"""
ResNet-50网络结构
"""
import tensorflow as tf


# 卷积层
def conv_op(input_tensor, scope_name, channel_output, training, bottle_neck, kh=3, kw=3, dh=1, dw=1, padding="SAME", activation=tf.nn.relu):
    channel_input = input_tensor.get_shape()[-1].value
    with tf.name_scope(scope_name) as scope:
        w = tf.get_variable(scope + "w", shape=[kh, kw, channel_input, channel_output], dtype=tf.float32,
                            initializer=tf.contrib.layers.xavier_initializer_conv2d())
        b = tf.get_variable(scope + "b", shape=[channel_output], dtype=tf.float32,
                            initializer=tf.constant_initializer(0.01))
        conv = tf.nn.conv2d(input_tensor, w, [1, dh, dw, 1], padding=padding)
        z = tf.nn.bias_add(conv, b)
        if bottle_neck:
            z = tf.layers.batch_normalization(z, trainable=training)
        if activation:
            z = activation(z)
        return z


# 最大池化层
def max_pool_op(input_tensor, scope_name, kh=2, kw=2, dh=2, dw=2,padding="SAME"):
    return tf.nn.max_pool(input_tensor, ksize=[1, kh, kw, 1], strides=[1, dh, dw, 1], padding=padding, name=scope_name)


# 全局池化层
def avg_pool_op(input_tensor, scope_name, kh=2, kw=2, dh=2, dw=2,padding="SAME"):
    return tf.nn.avg_pool(input_tensor, ksize=[1, kh, kw, 1], strides=[1, dh, dw, 1], padding=padding, name=scope_name)


# 全连接层
def fc_op(input_tensor, scope_name, channel_output, activation=tf.nn.relu):
    channel_input = input_tensor.get_shape()[-1].value
    with tf.name_scope(scope_name) as scope:
        w = tf.get_variable(scope + "w", shape=[channel_input,channel_output], dtype=tf.float32,
                            initializer=tf.contrib.layers.xavier_initializer())
        b = tf.get_variable(scope + "b", shape=[channel_output], dtype=tf.float32,
                            initializer=tf.constant_initializer(0.01))
        fc = tf.matmul(input_tensor, w) + b
    return fc


# bottle beck模块
def res_block_layers(input_tensor, scope_name, n_out_list, change_dimension=False, block_stride=1):
    if change_dimension:
        short_cut_conv = conv_op(input_tensor, scope_name + "_shortcut", n_out_list[1], training=True, bottle_neck=True,
                                 kh=1, kw=1, dh=block_stride, dw=block_stride, padding="SAME", activation=None)
    else:
        short_cut_conv = input_tensor
    block_conv_1 = conv_op(input_tensor, scope_name + "_bn_conv1", n_out_list[0], training=True, bottle_neck=True,
                           kh=1, kw=1, dh=block_stride, dw=block_stride, padding="SAME", activation=tf.nn.relu)
    block_conv_2 = conv_op(block_conv_1, scope_name + "_bn_conv2", n_out_list[0], training=True, bottle_neck=True,
                           kh=3, kw=3, dh=1, dw=1, padding="SAME", activation=tf.nn.relu)
    block_conv_3 = conv_op(block_conv_2, scope_name + "_bn_conv3", n_out_list[1], training=True, bottle_neck=True,
                           kh=1, kw=1, dh=1, dw=1, padding="SAME", activation=None)
    block_res = tf.add(short_cut_conv, block_conv_3)
    res = tf.nn.relu(block_res)
    return res


# ResNet-50
def resnet_50(input_tensor, num_classes, training=True, bottle_neck=True):
    # 第一模块
    conv1 = conv_op(input_tensor, "conv1", 64, training, bottle_neck, 3, 3, 1, 1)
    pool1 = max_pool_op(conv1, "pool1", kh=3, kw=3)
    # 第二模块
    block1_1 = res_block_layers(pool1, "block1_1", [64, 256], True, 1)
    block1_2 = res_block_layers(block1_1, "block1_2", [64, 256], False, 1)
    block1_3 = res_block_layers(block1_2, "block1_3", [64, 256], False, 1)
    # 第三模块
    block2_1 = res_block_layers(block1_3, "block2_1", [128, 512], True, 2)
    block2_2 = res_block_layers(block2_1, "block2_2", [128, 512], False, 1)
    block2_3 = res_block_layers(block2_2, "block2_3", [128, 512], False, 1)
    block2_4 = res_block_layers(block2_3, "block2_4", [128, 512], False, 1)
    # 第四模块
    block3_1 = res_block_layers(block2_4, "block3_1", [256, 1024], True, 2)
    block3_2 = res_block_layers(block3_1, "block3_2", [256, 1024], False, 1)
    block3_3 = res_block_layers(block3_2, "block3_3", [256, 1024], False, 1)
    block3_4 = res_block_layers(block3_3, "block3_4", [256, 1024], False, 1)
    block3_5 = res_block_layers(block3_4, "block3_5", [256, 1024], False, 1)
    block3_6 = res_block_layers(block3_5, "block3_6", [256, 1024], False, 1)
    # 第五模块
    block4_1 = res_block_layers(block3_6, "block4_1", [512, 2048], True, 2)
    block4_2 = res_block_layers(block4_1, "block4_2", [512, 2048], False, 1)
    block4_3 = res_block_layers(block4_2, "block4_3", [512, 2048], False, 1)
    pool2 = avg_pool_op(block4_3, "pool2", kh=7, kw=7, dh=1, dw=1, padding="SAME")
    # 展开pool2
    shape = pool2.get_shape()
    nodes = shape[1].value * shape[2].value * shape[3].value
    reshaped = tf.reshape(pool2, [-1, nodes], name="reshape")
    logits = fc_op(reshaped, "fc1", num_classes, activation=tf.nn.softmax)

    return logits
